package com.benromberg.junit.threadconductor;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.TimeoutException;

import org.junit.rules.TestRule;
import org.junit.runner.Description;
import org.junit.runners.model.Statement;

public class ThreadConductorRule implements TestRule {
	private Thread mainThread = Thread.currentThread();
	private List<ConductedThread> conductedThreads = new ArrayList<ConductedThread>();

	@Override
	public Statement apply(final Statement base, final Description description) {
		return new Statement() {
			@Override
			public void evaluate() throws Throwable {
				base.evaluate();
				for (ConductedThread conductedThread : conductedThreads) {
					if (conductedThread.isActive()) {
						conductedThread.waitForThread();
					}
					conductedThread.closeThread();
				}
				conductedThreads.clear();
			}
		};
	}

	public ConductedThread createThread() {
		ConductedThread conductedThread = new ConductedThread();
		conductedThreads.add(conductedThread);
		return conductedThread;
	}

	public class ConductedThread {
		private Thread thread;
		private Queue<Runnable> requestedRunnables = new LinkedList<Runnable>();
		private Throwable failure;
		private boolean active;
		private boolean exit;
		private boolean expectingInterrupt;
		private Map<Class<?>, Map<Integer, Object>> returnValueMap = new HashMap<Class<?>, Map<Integer, Object>>();

		private ConductedThread() {
			Runnable wrappedRunnable = new Runnable() {
				@Override
				public void run() {
					while (true) {
						synchronized (this) {
							try {
								wait();
								throw new RuntimeException(new TimeoutException());
							} catch (InterruptedException e) {
							}
						}
						while (!requestedRunnables.isEmpty()) {
							try {
								requestedRunnables.poll().run();
							} catch (Throwable t) {
								setFailure(t);
								break;
							}
						}
						active = false;
						if (expectingInterrupt) {
							mainThread.interrupt();
						}
						if (exit) {
							break;
						}
					}
				}
			};
			thread = new Thread(wrappedRunnable);
			thread.start();
		}

		public boolean isActive() {
			return active;
		}

		private void setFailure(Throwable failure) {
			this.failure = failure;
		}

		public <T> void setReturnValue(T returnValue) {
			setReturnValue(returnValue, 0);
		}

		public <T> void setReturnValue(T returnValue, int position) {
			if (returnValueMap.get(returnValue.getClass()) == null) {
				Map<Integer, Object> map = new HashMap<Integer, Object>();
				returnValueMap.put(returnValue.getClass(), map);
			}
			returnValueMap.get(returnValue.getClass()).put(Integer.valueOf(position), returnValue);
		}

		public <T> T getReturnValue(Class<T> returnValueClass) {
			return getReturnValue(returnValueClass, 0);
		}

		@SuppressWarnings("unchecked")
		public <T> T getReturnValue(Class<T> returnValueClass, int position) {
			return (T) returnValueMap.get(returnValueClass).get(Integer.valueOf(position));
		}

		private synchronized void waitForThread() throws Throwable {
			expectingInterrupt = true;
			while (active) {
				try {
					wait();
					throw new RuntimeException(new TimeoutException());
				} catch (InterruptedException e) {
				}
			}
			if (failure != null) {
				Throwable failureCache = failure;
				failure = null;
				throw failureCache;
			}
		}

		public void runAndWait(Runnable runnable) throws Throwable {
			runAndContinue(runnable);
			waitForThread();
		}

		public void runAndContinue(Runnable runnable) throws Throwable {
			requestedRunnables.add(runnable);
			active = true;
			thread.interrupt();
		}

		public void closeThread() {
			requestedRunnables.clear();
			exit = true;
			thread.interrupt();
		}
	}

}
